#include "md501.h"
#include <iomanip>
#include <assert.h>
#include <chrono>
#include <arm_neon.h>

using namespace std;
using namespace chrono;

/**
 * StringProcess: 将单个输入字符串转换成MD5计算所需的消息数组
 * @param input 输入
 * @param[out] n_byte 用于给调用者传递额外的返回值，即最终Byte数组的长度
 * @return Byte消息数组
 */
Byte *StringProcess(string input, int *n_byte)
{
    // 将输入的字符串转换为Byte为单位的数组
    Byte *blocks = (Byte *)input.c_str();
    int length = input.length();

    // 计算原始消息长度（以比特为单位）
    int bitLength = length * 8;

    // paddingBits: 原始消息需要的padding长度（以bit为单位）
    // 对于给定的消息，将其补齐至length%512==448为止
    // 需要注意的是，即便给定的消息满足length%512==448，也需要再pad 512bits
    int paddingBits = bitLength % 512;
    if (paddingBits > 448)
    {
        paddingBits += 512 - (paddingBits - 448);
    }
    else if (paddingBits < 448)
    {
        paddingBits = 448 - paddingBits;
    }
    else if (paddingBits == 448)
    {
        paddingBits = 512;
    }

    // 原始消息需要的padding长度（以Byte为单位）
    int paddingBytes = paddingBits / 8;
    // 创建最终的字节数组
    // length + paddingBytes + 8:
    // 1. length为原始消息的长度（bits）
    // 2. paddingBytes为原始消息需要的padding长度（Bytes）
    // 3. 在pad到length%512==448之后，需要额外附加64bits的原始消息长度，即8个bytes
    int paddedLength = length + paddingBytes + 8;
    Byte *paddedMessage = new Byte[paddedLength];

    // 复制原始消息
    memcpy(paddedMessage, blocks, length);

    // 添加填充字节。填充时，第一位为1，后面的所有位均为0。
    // 所以第一个byte是0x80
    paddedMessage[length] = 0x80;                            // 添加一个0x80字节
    memset(paddedMessage + length + 1, 0, paddingBytes - 1); // 填充0字节

    // 添加消息长度（64比特，小端格式）
    for (int i = 0; i < 8; ++i)
    {
        // 特别注意此处应当将bitLength转换为uint64_t
        // 这里的length是原始消息的长度
        paddedMessage[length + paddingBytes + i] = ((uint64_t)length * 8 >> (i * 8)) & 0xFF;
    }

    // 验证长度是否满足要求。此时长度应当是512bit的倍数
    int residual = 8 * paddedLength % 512;
    // assert(residual == 0);

    // 在填充+添加长度之后，消息被分为n_blocks个512bit的部分
    *n_byte = paddedLength;
    return paddedMessage;
}

/**
 * MD5Hash: 将单个输入字符串转换成MD5
 * @param input 输入
 * @param[out] state 用于给调用者传递额外的返回值，即最终的缓冲区，也就是MD5的结果
 * @return Byte消息数组
 */

void MD5Hash(const string inputs[4], uint32_t states[4][4]) // 支持 4 个输入
{
    // 1. 填充 4 个输入
    Byte *paddedMessages[4]; // 存储 4 个填充后的消息
    int messageLengths[4];   // 存储每个消息的长度

    for (int i = 0; i < 4; ++i)
    {
        paddedMessages[i] = StringProcess(inputs[i], &messageLengths[i]);
        assert(messageLengths[i] == messageLengths[0]); // 确保所有消息块长度一致
    }

    int n_blocks = messageLengths[0] / 64; // 每块 64 字节（512 位）

    // 2. 初始化 4 路 MD5 状态（向量化）
    uint32x4_t state_a = vdupq_n_u32(0x67452301); // A
    uint32x4_t state_b = vdupq_n_u32(0xefcdab89); // B
    uint32x4_t state_c = vdupq_n_u32(0x98badcfe); // C
    uint32x4_t state_d = vdupq_n_u32(0x10325476); // D

    // 3. 分块处理
    for (int block = 0; block < n_blocks; ++block)
    {
        uint32x4_t x[16]; // 存储当前块的 16 个 32 位向量

        // 加载当前块的 16 个 32 位向量
        for (int i = 0; i < 16; ++i)
        {
            uint32_t x_scalar[4];
            for (int j = 0; j < 4; ++j)
            {
                Byte *block_ptr = paddedMessages[j] + block * 64;
                x_scalar[j] = (block_ptr[4 * i]) |
                              (block_ptr[4 * i + 1] << 8) |
                              (block_ptr[4 * i + 2] << 16) |
                              (block_ptr[4 * i + 3] << 24);
            }
            x[i] = vld1q_u32(x_scalar); // 加载到 NEON 向量
        }

        // 保存当前状态
        uint32x4_t a = state_a, b = state_b, c = state_c, d = state_d;

        /* Round 1 */
        FF_SIMD(a, b, c, d, x[0], 7, vdupq_n_u32(0xd76aa478));
        FF_SIMD(d, a, b, c, x[1], 12, vdupq_n_u32(0xe8c7b756));
        FF_SIMD(c, d, a, b, x[2], 17, vdupq_n_u32(0x242070db));
        FF_SIMD(b, c, d, a, x[3], 22, vdupq_n_u32(0xc1bdceee));

        FF_SIMD(a, b, c, d, x[4], 7, vdupq_n_u32(0xf57c0faf));
        FF_SIMD(d, a, b, c, x[5], 12, vdupq_n_u32(0x4787c62a));
        FF_SIMD(c, d, a, b, x[6], 17, vdupq_n_u32(0xa8304613));
        FF_SIMD(b, c, d, a, x[7], 22, vdupq_n_u32(0xfd469501));

        FF_SIMD(a, b, c, d, x[8], 7, vdupq_n_u32(0x698098d8));
        FF_SIMD(d, a, b, c, x[9], 12, vdupq_n_u32(0x8b44f7af));
        FF_SIMD(c, d, a, b, x[10], 17, vdupq_n_u32(0xffff5bb1));
        FF_SIMD(b, c, d, a, x[11], 22, vdupq_n_u32(0x895cd7be));

        FF_SIMD(a, b, c, d, x[12], 7, vdupq_n_u32(0x6b901122));
        FF_SIMD(d, a, b, c, x[13], 12, vdupq_n_u32(0xfd987193));
        FF_SIMD(c, d, a, b, x[14], 17, vdupq_n_u32(0xa679438e));
        FF_SIMD(b, c, d, a, x[15], 22, vdupq_n_u32(0x49b40821));

        /* Round 2 */
        GG_SIMD(a, b, c, d, x[1], 5, vdupq_n_u32(0xf61e2562));
        GG_SIMD(d, a, b, c, x[6], 9, vdupq_n_u32(0xc040b340));
        GG_SIMD(c, d, a, b, x[11], 14, vdupq_n_u32(0x265e5a51));
        GG_SIMD(b, c, d, a, x[0], 20, vdupq_n_u32(0xe9b6c7aa));

        GG_SIMD(a, b, c, d, x[5], 5, vdupq_n_u32(0xd62f105d));
        GG_SIMD(d, a, b, c, x[10], 9, vdupq_n_u32(0x02441453));
        GG_SIMD(c, d, a, b, x[15], 14, vdupq_n_u32(0xd8a1e681));
        GG_SIMD(b, c, d, a, x[4], 20, vdupq_n_u32(0xe7d3fbc8));

        GG_SIMD(a, b, c, d, x[9], 5, vdupq_n_u32(0x21e1cde6));
        GG_SIMD(d, a, b, c, x[14], 9, vdupq_n_u32(0xc33707d6));
        GG_SIMD(c, d, a, b, x[3], 14, vdupq_n_u32(0xf4d50d87));
        GG_SIMD(b, c, d, a, x[8], 20, vdupq_n_u32(0x455a14ed));

        GG_SIMD(a, b, c, d, x[13], 5, vdupq_n_u32(0xa9e3e905));
        GG_SIMD(d, a, b, c, x[2], 9, vdupq_n_u32(0xfcefa3f8));
        GG_SIMD(c, d, a, b, x[7], 14, vdupq_n_u32(0x676f02d9));
        GG_SIMD(b, c, d, a, x[12], 20, vdupq_n_u32(0x8d2a4c8a));

        /* Round 3 */
        HH_SIMD(a, b, c, d, x[5], 4, vdupq_n_u32(0xfffa3942));
        HH_SIMD(d, a, b, c, x[8], 11, vdupq_n_u32(0x8771f681));
        HH_SIMD(c, d, a, b, x[11], 16, vdupq_n_u32(0x6d9d6122));
        HH_SIMD(b, c, d, a, x[14], 23, vdupq_n_u32(0xfde5380c));

        HH_SIMD(a, b, c, d, x[1], 4, vdupq_n_u32(0xa4beea44));
        HH_SIMD(d, a, b, c, x[4], 11, vdupq_n_u32(0x4bdecfa9));
        HH_SIMD(c, d, a, b, x[7], 16, vdupq_n_u32(0xf6bb4b60));
        HH_SIMD(b, c, d, a, x[10], 23, vdupq_n_u32(0xbebfbc70));

        HH_SIMD(a, b, c, d, x[13], 4, vdupq_n_u32(0x289b7ec6));
        HH_SIMD(d, a, b, c, x[0], 11, vdupq_n_u32(0xeaa127fa));
        HH_SIMD(c, d, a, b, x[3], 16, vdupq_n_u32(0xd4ef3085));
        HH_SIMD(b, c, d, a, x[6], 23, vdupq_n_u32(0x04881d05));

        HH_SIMD(a, b, c, d, x[9], 4, vdupq_n_u32(0xd9d4d039));
        HH_SIMD(d, a, b, c, x[12], 11, vdupq_n_u32(0xe6db99e5));
        HH_SIMD(c, d, a, b, x[15], 16, vdupq_n_u32(0x1fa27cf8));
        HH_SIMD(b, c, d, a, x[2], 23, vdupq_n_u32(0xc4ac5665));

        /* Round 4 */
        II_SIMD(a, b, c, d, x[0], 6, vdupq_n_u32(0xf4292244));
        II_SIMD(d, a, b, c, x[7], 10, vdupq_n_u32(0x432aff97));
        II_SIMD(c, d, a, b, x[14], 15, vdupq_n_u32(0xab9423a7));
        II_SIMD(b, c, d, a, x[5], 21, vdupq_n_u32(0xfc93a039));

        II_SIMD(a, b, c, d, x[12], 6, vdupq_n_u32(0x655b59c3));
        II_SIMD(d, a, b, c, x[3], 10, vdupq_n_u32(0x8f0ccc92));
        II_SIMD(c, d, a, b, x[10], 15, vdupq_n_u32(0xffeff47d));
        II_SIMD(b, c, d, a, x[1], 21, vdupq_n_u32(0x85845dd1));

        II_SIMD(a, b, c, d, x[8], 6, vdupq_n_u32(0x6fa87e4f));
        II_SIMD(d, a, b, c, x[15], 10, vdupq_n_u32(0xfe2ce6e0));
        II_SIMD(c, d, a, b, x[6], 15, vdupq_n_u32(0xa3014314));
        II_SIMD(b, c, d, a, x[13], 21, vdupq_n_u32(0x4e0811a1));

        II_SIMD(a, b, c, d, x[4], 6, vdupq_n_u32(0xf7537e82));
        II_SIMD(d, a, b, c, x[11], 10, vdupq_n_u32(0xbd3af235));
        II_SIMD(c, d, a, b, x[2], 15, vdupq_n_u32(0x2ad7d2bb));
        II_SIMD(b, c, d, a, x[9], 21, vdupq_n_u32(0xeb86d391));

        // 更新全局状态
        state_a = vaddq_u32(state_a, a);
        state_b = vaddq_u32(state_b, b);
        state_c = vaddq_u32(state_c, c);
        state_d = vaddq_u32(state_d, d);
    }

    // // 4. 存储结果并调整字节序
    // vst1q_u32(states[0], state_a);
    // vst1q_u32(states[1], state_b);
    // vst1q_u32(states[2], state_c);
    // vst1q_u32(states[3], state_d);

    // for (int idx = 0; idx < 4; ++idx)
    // {
    //     for (int i = 0; i < 4; ++i)
    //     {
    //         uint32_t value = states[idx][i];
    //         states[idx][i] = ((value & 0xff) << 24) |
    //                          ((value & 0xff00) << 8) |
    //                          ((value & 0xff0000) >> 8) |
    //                          ((value & 0xff000000) >> 24);
    //     }
    // }
    // 修改存储结果部分
uint32_t temp_a[4], temp_b[4], temp_c[4], temp_d[4];
vst1q_u32(temp_a, state_a);
vst1q_u32(temp_b, state_b);
vst1q_u32(temp_c, state_c);
vst1q_u32(temp_d, state_d);

// 转置矩阵
for (int i = 0; i < 4; ++i) {
    states[i][0] = temp_a[i];
    states[i][1] = temp_b[i];
    states[i][2] = temp_c[i];
    states[i][3] = temp_d[i];
}

// 然后再进行字节序调整
for (int i = 0; i < 4; ++i) {
    for (int j = 0; j < 4; ++j) {
        uint32_t value = states[i][j];
        states[i][j] = ((value & 0xff) << 24) |
                       ((value & 0xff00) << 8) |
                       ((value & 0xff0000) >> 8) |
                       ((value & 0xff000000) >> 24);
    }
}

    // 释放动态分配的内存
    for (int i = 0; i < 4; ++i)
        delete[] paddedMessages[i];
}